{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CqWhSlfDqmns"
   },
   "source": [
    "<a id='prerequisits'></a>\n",
    "\n",
    "# Prerequisits\n",
    "\n",
    "This section installs `gradslam` (if not already installed), imports the necessary packages for the tutorial, and downloads 'lr kt1' (the first trajectory) of [ICL-NUIM dataset](https://www.doc.ic.ac.uk/~ahanda/VaFRIC/iclnuim.html) and structures it as below:\n",
    "```\n",
    "ICL\n",
    "    living_room_traj1_frei_png\n",
    "        depth/    rgb/    associations.txt    livingRoom1.gt.freiburg    livingRoom1n.gt.sim\n",
    "```\n",
    "\n",
    "\n",
    "We set the ICL path variable: `icl_path='ICL/'`. The ICL data is loaded into the following variables: <br>\n",
    "\n",
    "* `colors`: of shape (batch_size, sequence_length, height, width, 3) <br>\n",
    "* `depths`: of shape (batch_size, sequence_length, height, width, 1) <br>\n",
    "* `intrinsics`: of shape (batch_size, 1, 4, 4) <br>\n",
    "* `poses`: of shape (batch_size, sequence_length, 4, 4) <br>\n",
    "\n",
    "Finally `RGBDImages` is created from the ICL data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9opFP8sDqmnt"
   },
   "outputs": [],
   "source": [
    "# install gradslam (if not installed)\n",
    "try:\n",
    "    import gradslam as gs\n",
    "except ImportError:\n",
    "    print(\"Installing gradslam...\")\n",
    "    !pip install 'git+https://github.com/gradslam/gradslam.git' -q\n",
    "    print('Installed')\n",
    "\n",
    "# import necessary packages\n",
    "import gradslam as gs\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "from gradslam import Pointclouds, RGBDImages\n",
    "from gradslam.datasets import ICL\n",
    "from gradslam.slam import PointFusion\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# download 'lr kt1' of ICL dataset\n",
    "if not os.path.isdir('ICL'):\n",
    "    os.mkdir('ICL')\n",
    "if not os.path.isdir('ICL/living_room_traj1_frei_png'):\n",
    "    print('Downloading ICL/living_room_traj1_frei_png dataset...')\n",
    "    os.mkdir('ICL/living_room_traj1_frei_png')\n",
    "    !wget http://www.doc.ic.ac.uk/~ahanda/living_room_traj1_frei_png.tar.gz -P ICL/living_room_traj1_frei_png/ -q\n",
    "    !tar -xzf ICL/living_room_traj1_frei_png/living_room_traj1_frei_png.tar.gz -C ICL/living_room_traj1_frei_png/\n",
    "    !rm ICL/living_room_traj1_frei_png/living_room_traj1_frei_png.tar.gz\n",
    "    !wget https://www.doc.ic.ac.uk/~ahanda/VaFRIC/livingRoom1n.gt.sim -P ICL/living_room_traj1_frei_png/ -q\n",
    "    print('Downloaded.')\n",
    "icl_path = 'ICL/'\n",
    "\n",
    "# load dataset\n",
    "dataset = ICL(icl_path, seqlen=8, height=240, width=320)\n",
    "loader = DataLoader(dataset=dataset, batch_size=2)\n",
    "colors, depths, intrinsics, poses, *_ = next(iter(loader))\n",
    "\n",
    "# create rgbdimages object\n",
    "rgbdimages = RGBDImages(colors, depths, intrinsics, poses)\n",
    "rgbdimages.plotly(0).update_layout(autosize=False, height=600, width=400).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z7KZsY2lqLe1"
   },
   "source": [
    "# Instantiation\n",
    "\n",
    "> **_NOTE:_**  Make sure to have ran the [prerequisits](#Prerequisits) section before running this section.\n",
    "\n",
    "The `Pointclouds` structures aims to contain batched pointclouds and allow for batched operations on pointclouds.\n",
    "\n",
    "A `Pointclouds` object can be initialized from points coordinates, point normals, point colors and (optionally) point features. These attributes can be passed in one of the following representations:\n",
    "\n",
    "- List (of `torch.Tensor` objects): Store points of each pointcloud of shape $(𝑁_𝑏, 3)$ in a list of $𝐵$ `torch.Tensor` objects.\n",
    "\n",
    "- Padded: Store all points in a $(𝐵, N, 3)$ tensor.\n",
    "\n",
    "`Pointclouds` can also be instantiated from `RGBDImages`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VekYeIKkqbsF",
    "outputId": "7eac7a2b-232d-4f2d-dc3b-568dbc00b0b2"
   },
   "outputs": [],
   "source": [
    "from gradslam.structures.utils import pointclouds_from_rgbdimages\n",
    "\n",
    "# instantiate empty Pointclouds object\n",
    "pointclouds = Pointclouds()\n",
    "print(pointclouds.has_points)  # False\n",
    "print('---')\n",
    "\n",
    "# instantiation from list of tensors of points\n",
    "pointclouds = Pointclouds(points=[torch.rand(4, 3), torch.rand(2, 3), torch.rand(1, 3)],\n",
    "                          normals=[torch.rand(4, 3), torch.rand(2, 3), torch.rand(1, 3)],\n",
    "                          colors=[torch.rand(4, 3), torch.rand(2, 3), torch.rand(1, 3)])\n",
    "print(pointclouds.num_points_per_pointcloud)  # tensor([4, 2, 1])\n",
    "print('---')\n",
    "\n",
    "# instantiation from tensor\n",
    "pointclouds = Pointclouds(points=torch.rand(3, 4, 3),\n",
    "                          normals=torch.rand(3, 4, 3),\n",
    "                          colors=torch.rand(3, 4, 3))\n",
    "print(pointclouds.num_points_per_pointcloud)  # tensor([4, 4, 4])\n",
    "print('---')\n",
    "\n",
    "# instantiate with features\n",
    "# features can have any number of dimensions\n",
    "pointclouds = Pointclouds(points=torch.rand(3, 4, 3),\n",
    "                          normals=torch.rand(3, 4, 3),\n",
    "                          colors=torch.rand(3, 4, 3),\n",
    "                          features=torch.rand(3, 4, 10))\n",
    "print(pointclouds.has_features)  # True\n",
    "print('---')\n",
    "\n",
    "# instantiate from RGBDImages with sequence length of 1\n",
    "rgbdimages1 = rgbdimages[:, 0]\n",
    "pointclouds = pointclouds_from_rgbdimages(rgbdimages1, filter_missing_depths=False)\n",
    "print(rgbdimages1.shape)  # (2, 1, 240, 320)\n",
    "print(pointclouds.num_points_per_pointcloud)  # tensor([76800, 76800])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w97HUr2PhH36"
   },
   "source": [
    "# List and padded internal representations\n",
    "\n",
    "[Similar to PyTorch3d](https://pytorch3d.org/docs/batching), our `Pointclouds` structure suppors a list representation and a padded representation internally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wcCwXFOmhUVU",
    "outputId": "3aa1b03b-0aab-42af-8b45-147cab0af3c8"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "# instantiation from list of tensors of points\n",
    "pointclouds = Pointclouds(points=[torch.rand(2, 3), torch.rand(1, 3)],\n",
    "                          normals=[torch.rand(2, 3), torch.rand(1, 3)],\n",
    "                          colors=[torch.rand(2, 3), torch.rand(1, 3)],\n",
    "                          features=[torch.rand(2, 10), torch.rand(1, 10)])\n",
    "print(pointclouds.num_points_per_pointcloud)  # tensor([5, 2, 8])\n",
    "print('---')\n",
    "\n",
    "# List representation\n",
    "for points in pointclouds.points_list:\n",
    "    print(points)\n",
    "# tensor([[0.4963, 0.7682, 0.0885],\n",
    "#         [0.1320, 0.3074, 0.6341]])\n",
    "# tensor([[0.4901, 0.8964, 0.4556]])\n",
    "print('---')\n",
    "\n",
    "# Padded representation\n",
    "print(pointclouds.points_padded)\n",
    "# tensor([[[0.4963, 0.7682, 0.0885],\n",
    "#          [0.1320, 0.3074, 0.6341]],\n",
    "\n",
    "#         [[0.4901, 0.8964, 0.4556],\n",
    "#          [0.0000, 0.0000, 0.0000]]])\n",
    "print('---')\n",
    "\n",
    "# Padded representation shapes\n",
    "print(pointclouds.points_padded.shape)  # torch.Size([2, 2, 3])\n",
    "print(pointclouds.normals_padded.shape)  # torch.Size([2, 2, 3])\n",
    "print(pointclouds.colors_padded.shape)  # torch.Size([2, 2, 3])\n",
    "print(pointclouds.features_padded.shape)  # torch.Size([2, 2, 10])\n",
    "print('---')\n",
    "\n",
    "# List representation shapes\n",
    "print([p.shape for p in pointclouds.points_list])  # [torch.Size([2, 3]), torch.Size([1, 3])]\n",
    "print([n.shape for n in pointclouds.normals_list])  # [torch.Size([2, 3]), torch.Size([1, 3])]\n",
    "print([c.shape for c in pointclouds.colors_list])  # [torch.Size([2, 3]), torch.Size([1, 3])]\n",
    "print([f.shape for f in pointclouds.features_list])  # [torch.Size([2, 10]), torch.Size([1, 10])]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4aAK9HbIqmnv"
   },
   "source": [
    "# Indexing and slicing\n",
    "\n",
    "Basic indexing and slicing of `Pointclouds` over the first (batch) dimension is supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xlqE42usqmnv",
    "outputId": "01cb2207-81e7-4bcd-90f1-49d933675aa7"
   },
   "outputs": [],
   "source": [
    "# initalize Pointclouds\n",
    "pointclouds = Pointclouds(points=torch.rand(3, 4, 3),\n",
    "                          normals=torch.rand(3, 4, 3),\n",
    "                          colors=torch.rand(3, 4, 3))\n",
    "print(len(pointclouds))  # 3\n",
    "print('---')\n",
    "\n",
    "# indexing\n",
    "pointclouds1 = pointclouds[0]\n",
    "print(len(pointclouds1))  # 1\n",
    "print('---')\n",
    "\n",
    "# slicing\n",
    "pointclouds2 = pointclouds[:2]\n",
    "print(len(pointclouds2))  # 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hPIkUUNqqmnx"
   },
   "source": [
    "# Translations, rotations and transformations\n",
    "\n",
    "`Pointclouds` supports batch mode geometric operations such as translations, rotations, and transformations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "warZULSpqmny",
    "outputId": "51b0f8bf-d2b2-4426-f670-925f9ae9427b",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import plotly.graph_objects as go\n",
    "torch.manual_seed(0)\n",
    "\n",
    "def custom_plotly_viz(pointclouds, title):\n",
    "    fig = go.Figure(pointclouds.plotly(0, as_figure=False, point_size=15))\n",
    "    fig.update_layout(title=title, autosize=False, height=400, width=400)\n",
    "    fig.show()\n",
    "\n",
    "# initalize Pointclouds\n",
    "pointclouds = Pointclouds(points=torch.rand(2, 5, 3),\n",
    "                          normals=torch.rand(2, 5, 3),\n",
    "                          colors=torch.rand(2, 5, 3))\n",
    "\n",
    "# translate\n",
    "pointclouds1 = pointclouds + 10\n",
    "\n",
    "# scale\n",
    "pointclouds2 = pointclouds * 100\n",
    "\n",
    "# rotate (Bx3x3 rotation)\n",
    "rmat = torch.tensor(\n",
    "    [\n",
    "     [(3 ** 0.5) / 2, -0.5, 0],\n",
    "     [0.5, (3 ** 0.5) / 2, 0],\n",
    "     [0, 0, 1],\n",
    "     ]\n",
    "     )\n",
    "pointclouds3 = pointclouds.rotate(rmat)\n",
    "\n",
    "# transform (Bx4x4 transformation)\n",
    "mat = torch.tensor(\n",
    "    [\n",
    "     [(3 ** 0.5) / 2, -0.5, 0, 20],\n",
    "     [0.5, (3 ** 0.5) / 2, 0, 20],\n",
    "     [0, 0, 1, 20],\n",
    "     [0, 0, 0, 1],\n",
    "     ]\n",
    "     )\n",
    "pointclouds4 = pointclouds.transform(mat)\n",
    "\n",
    "# visualizations\n",
    "custom_plotly_viz(pointclouds, \"pointclouds[0]\")\n",
    "custom_plotly_viz(pointclouds1, \"pointclouds[0] + 10\")\n",
    "custom_plotly_viz(pointclouds2, \"pointclouds[0] * 100\")\n",
    "custom_plotly_viz(pointclouds3, \"pointclouds[0] rotated 30 deg about z-axis\")\n",
    "custom_plotly_viz(pointclouds4, \"pointclouds[0] rigid transformation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r7eUMXb5zoZz"
   },
   "source": [
    "# Pinhole camera projection\n",
    "`Pointclouds` can be projected onto a 2-d plane given the intrinsics matrix using the `Pointclouds.pinhole_projection(intrinsics)` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 817
    },
    "id": "TjCsNicsznh1",
    "outputId": "1e2e5ede-3f53-4019-c9aa-580e32d763b4"
   },
   "outputs": [],
   "source": [
    "import plotly.graph_objects as go\n",
    "torch.manual_seed(0)\n",
    "\n",
    "def custom_plotly_viz(pointclouds, title):\n",
    "    fig = go.Figure(pointclouds.plotly(0, as_figure=False, point_size=15))\n",
    "    fig.update_layout(title=title, autosize=False, height=400, width=400)\n",
    "    fig.show()\n",
    "\n",
    "# initalize Pointclouds\n",
    "pointclouds = Pointclouds(points=torch.rand(2, 5, 3),\n",
    "                          normals=torch.rand(2, 5, 3),\n",
    "                          colors=torch.rand(2, 5, 3))\n",
    "\n",
    "# pinhole projection\n",
    "fx, fy = 1., 1.\n",
    "cx, cy = 0.5, 0.5\n",
    "intrinsics = torch.tensor(\n",
    "    [\n",
    "     [fx, 0, cx, 0],\n",
    "     [0, fy, cy, 0],\n",
    "     [0, 0, 1, 0],\n",
    "     [0, 0, 0, 1]\n",
    "    ]\n",
    ")\n",
    "pointclouds1 = pointclouds.pinhole_projection(intrinsics)\n",
    "\n",
    "# visualizations\n",
    "custom_plotly_viz(pointclouds, \"pointclouds[0]\")\n",
    "custom_plotly_viz(pointclouds1, \"pointclouds[0] projection\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qfpvR58Sqmn0"
   },
   "source": [
    "# Transfer between GPU/CPU\n",
    "\n",
    "`Pointclouds` support easy transfer between CPU and GPU. This operation transfers all tensors in the `Pointclouds` objects between CPU/GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EbxiaFwSuNvv",
    "outputId": "43c8c9e6-703e-44b2-84a5-b61a56a64fcd"
   },
   "outputs": [],
   "source": [
    "# initalize Pointclouds\n",
    "pointclouds = Pointclouds(points=torch.rand(2, 5, 3),\n",
    "                          normals=torch.rand(2, 5, 3),\n",
    "                          colors=torch.rand(2, 5, 3))\n",
    "\n",
    "# transfer to GPU\n",
    "if torch.cuda.is_available():\n",
    "    pointclouds = pointclouds.to(\"cuda\")\n",
    "    pointclouds = pointclouds.cuda()  # equivalent to pointclouds.to(\"cuda\")\n",
    "    print(pointclouds.points_padded.device)  # \"cuda:0\"\n",
    "    print('---')\n",
    "\n",
    "# transfer to CPU\n",
    "pointclouds = pointclouds.to(\"cpu\")\n",
    "pointclouds = pointclouds.cpu()  # equivalent to pointclouds.to(\"cpu\")\n",
    "print(pointclouds.points_padded.device)  # \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yrG8sexkuRBK"
   },
   "source": [
    "# Detach and clone tensors\n",
    "\n",
    "`Pointclouds.detach` returns a new `Pointclouds` object such that all internal tensors of the new object do not require grad. `Pointclouds.clone()` returns a new `Pointclouds` object such that all the internal tensors are cloned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q9vkSwS1uWNc",
    "outputId": "6450824a-542f-40d4-e25a-c4188661e325"
   },
   "outputs": [],
   "source": [
    "# initalize Pointclouds\n",
    "pointclouds = Pointclouds(points=torch.rand((2, 5, 3), requires_grad=True),\n",
    "                          normals=torch.rand((2, 5, 3), requires_grad=True),\n",
    "                          colors=torch.rand((2, 5, 3), requires_grad=True))\n",
    "\n",
    "# clone\n",
    "pointclouds1 = pointclouds.clone()\n",
    "print(torch.allclose(pointclouds1.points_padded, pointclouds.points_padded))  # True\n",
    "print(pointclouds1.points_padded is pointclouds.points_padded)  # False\n",
    "print('---')\n",
    "\n",
    "# detach\n",
    "pointclouds2 = pointclouds.detach()\n",
    "print(pointclouds.points_padded.requires_grad)  # True\n",
    "print(pointclouds2.points_padded.requires_grad)  # False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lhHwfuA2mj4j"
   },
   "source": [
    "# Visualization\n",
    "\n",
    "`Pointclouds` can quickly and easily be visualized with either the `.plotly(batch_index)` method or the the `.open3d(batch_index)` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 542
    },
    "id": "7iUau5mvmqFV",
    "outputId": "045af1ec-4356-4747-cd96-9bda079d36a5"
   },
   "outputs": [],
   "source": [
    "import open3d as o3d\n",
    "torch.manual_seed(0)\n",
    "\n",
    "# initalize Pointclouds\n",
    "pointclouds = Pointclouds(points=torch.rand((2, 5, 3), requires_grad=True),\n",
    "                          normals=torch.rand((2, 5, 3), requires_grad=True),\n",
    "                          colors=torch.rand((2, 5, 3), requires_grad=True))\n",
    "\n",
    "# plotly visualization\n",
    "pointclouds.plotly(0, point_size=20).show()\n",
    "\n",
    "# open3d visualization (does not work with Google Colab)\n",
    "# o3d.visualization.draw_geometries([pointclouds.open3d(0)])"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "pointclouds_tutorial.ipynb",
   "provenance": []
  },
  "jupytext": {
   "text_representation": {
    "extension": ".py",
    "format_name": "light",
    "format_version": "1.5",
    "jupytext_version": "1.6.0"
   }
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
